#!/usr/bin/env python
# -*- coding:utf-8 -*-
'''
@File   :   make_train_mgf.py
@Author :   Song
@Time   :   2022/5/16 10:55
@Contact:   songjian@westlake.edu.cn
@intro  : 
'''
import time
import numpy as np
import pandas as pd
import multiprocessing as mp
from pathlib import Path
import pyteomics.mzml
import pyteomics.mgf
import click
import difflib

class mzml_Reader():

    def __init__(self, fpath):
        self.mz = None
        self.all_rt = np.array([])  # 表示每一次采集的rt
        self.TimeUnit = None

        # type
        fpath = str(fpath)
        self.suffix = fpath.split('.')[-1].lower()

        self.load(fpath)
        self.init()

    def size(self):
        return len(self.mz)

    def load(self, file_path):
        assert self.suffix == 'mzml', 'Only .mzML is supported!'
        self.mz = pyteomics.mzml.MzML(file_path, use_index=True)

    def init(self):
        self.get_time_unit()
        self.__load_to_memory()

    def get_mz_suffix(self):
        return self.suffix

    def get_time_unit(self):
        if self.suffix == 'mzxml':
            if self.mz[self.size() - 1]['retentionTime'] < 3 * 60:
                self.TimeUnit = 'minute'
            else:
                self.TimeUnit = 'second'
        elif self.suffix == 'mzml':
            if self.mz[self.size() - 1]['scanList']['scan'][0]['scan start time'] < 5 * 60:
                self.TimeUnit = 'minute'
            else:
                self.TimeUnit = 'second'

    def process_worker(self, idx_start, idx_end): # 作为进程worker不能私有
        rts, levels, peaks_mz, peaks_intensity, hcds = [], [], [], [], []
        pr_mz_v, pr_charge_v = [], []
        for idx in range(idx_start, idx_end):
            scan = self.mz[idx]
            rts.append(np.float32(scan['scanList']['scan'][0]['scan start time']))
            level = np.int8(scan['ms level'])
            levels.append(level)
            if level == 2:
                scan_pr = scan['precursorList']['precursor'][-1]
                hcds.append(np.float32(scan_pr['activation']['collision energy']))
                pr = scan_pr['selectedIonList']['selectedIon'][-1]
                pr_mz_v.append(np.float32(pr['selected ion m/z']))
                pr_charge_v.append(np.int8(pr['charge state']))
            else:
                hcds.append(-1)
                pr_mz_v.append(-1)
                pr_charge_v.append(-1)
            peaks_mz.append(scan['m/z array'].astype(np.float32))
            peaks_intensity.append(scan['intensity array'].astype(np.float32))

        return {'rt': rts, 'level': levels, 'hcd': hcds,
                'mz': peaks_mz, 'intensity': peaks_intensity,
                'pr_mz': pr_mz_v, 'pr_charge': pr_charge_v}

    def __load_to_memory(self):
        self.all_mz, self.all_intensity, self.all_levels, self.all_rt = [], [], [], []
        self.all_hcd = []
        self.all_pr_mz, self.all_pr_charge = [], []

        # multiprocess load data
        cpu_num = mp.cpu_count()
        process_num = int(cpu_num / 2)  # default cores / 2
        process_num = 8 if process_num >= 8 else process_num
        pool = mp.Pool(process_num)
        slices = np.ceil(np.linspace(0, len(self.mz), process_num+1)).astype(int)

        results = [pool.apply_async(self.process_worker, args=(slices[i], slices[i+1])) for i in range(process_num)]
        results = [r.get() for r in results] # get同步进程。result是四个字典
        pool.close()
        pool.join()

        for result in results:
            self.all_rt.extend(result['rt'])
            self.all_levels.extend(result['level'])
            self.all_hcd.extend(result['hcd'])
            self.all_mz.extend(result['mz'])
            self.all_intensity.extend(result['intensity'])
            self.all_pr_mz.extend(result['pr_mz'])
            self.all_pr_charge.extend(result['pr_charge'])

        self.all_rt = np.array(self.all_rt)
        self.all_levels = np.array(self.all_levels)
        self.all_hcd = np.array(self.all_hcd)
        self.all_mz = np.array(self.all_mz, dtype=object)
        self.all_intensity = np.array(self.all_intensity, dtype=object)
        self.all_scan_len = np.array(list(map(len, self.all_mz)))
        self.all_pr_mz = np.array(self.all_pr_mz)
        self.all_pr_charge = np.array(self.all_pr_charge)

        # 时间处理
        if self.TimeUnit == 'minute':
            self.all_rt = self.all_rt * 60.


def get_sample_idx(df, mzml):
    df = df[df['hcd'] == 35]
    df = df[df['Charge'] < 5]

    idx_maxquant = set(df['Scan number'].to_numpy() - 1)
    idx_msms = set(np.where((mzml.all_levels == 2) &
                            (mzml.all_hcd == 35) &
                            (mzml.all_pr_charge < 5))[0])
    # assert idx_maxquant.issubset(idx_msms)
    idx_neg = idx_msms - idx_maxquant

    # 正样本
    df = df[df['Score'] > 60]
    df = df[df['Length'] < 25]
    idx_pos = set(df['Scan number'].to_numpy() - 1)
    assert len(idx_pos & idx_neg) == 0

    idx_pos = np.array(list(idx_pos))
    idx_neg = np.array(list(idx_neg))
    # assert len(idx_pos) > 1000
    # idx_pos = np.random.choice(idx_pos, size=1000, replace=False)
    # idx_neg = np.random.choice(idx_neg, size=1000, replace=False)

    return idx_pos, idx_neg


@click.command()
@click.option('--dir_mzml', required=True, help='Specify the path of train_mzml')
@click.option('--dir_msms', required=True, help='Specify the path of train_msms')
@click.option('--dir_npz', required=True, help='Specify the output path of train_mgf')
def main(dir_mzml, dir_msms, dir_npz):
    dir_npz = Path(dir_npz)
    dir_mzml = Path(dir_mzml).absolute()
    dir_msms = Path(dir_msms).absolute()

    # 检查mzml和msms文件名一致
    dir_mzml_all = [_.name for _ in dir_mzml.iterdir()]
    dir_msms_all = [_.name for _ in dir_msms.iterdir()]
    dir_msms_all_match = pd.Series(dir_msms_all).replace('-', '_', regex=True)
    assert len(dir_mzml_all) == len(dir_msms_all), '两个输入文件夹数目不一致！'

    # 读取mzml和成对的msms
    use_cols = ['Score', 'Length', 'Modifications', 'Modified sequence',
                'm/z', 'Charge', 'Scan number', 'Retention time']
    pr_mz_v, pr_charge_v = [], []
    scan_x_v, scan_y_v, scan_len_v = [], [], []
    label_v = []
    total = len(list(dir_mzml.iterdir()))

    for file_idx, mzml_d in enumerate(dir_mzml.iterdir()):
        if file_idx == 0:
            continue
        x = mzml_d.name
        x = x[x.find('T') : -5].replace('-', '_') # 从T开始，到.mzML结束
        match_idx = np.where(dir_msms_all_match.str.contains(x))[0]
        assert len(match_idx) == 1
        fname_msms = dir_msms_all[int(match_idx)]
        msms_d = dir_msms/fname_msms

        print(f'{file_idx+1}/{total}, {mzml_d.name}, {msms_d.name}')

        mzml = mzml_Reader(mzml_d)

        df = pd.read_csv(msms_d, sep='\t', usecols=use_cols)
        # 添加谱图对应的HCD
        assert (mzml.all_levels[df['Scan number'] - 1] == 2).all()
        df['hcd'] = mzml.all_hcd[df['Scan number'] - 1]
        idx_pos, idx_neg = get_sample_idx(df, mzml) # idx_pos表示正样本，idx_maxquant表示有打分的谱图
        idx_pos_neg = np.concatenate([idx_pos, idx_neg])
        print(f'pos num: {len(idx_pos)}, '
              f'neg num: {len(idx_neg)}')

        # 提取
        pr_mz_v.append(mzml.all_pr_mz[idx_pos_neg])
        pr_charge_v.append(mzml.all_pr_charge[idx_pos_neg])
        scan_x_v.extend(mzml.all_mz[idx_pos_neg])
        scan_y_v.extend(mzml.all_intensity[idx_pos_neg])
        scan_len_v.append(mzml.all_scan_len[idx_pos_neg])
        label_v.append([1] * len(idx_pos))
        label_v.append([0] * len(idx_neg))

    pr_mz = np.concatenate(pr_mz_v, dtype=np.float32)
    pr_charge = np.concatenate(pr_charge_v, dtype=np.int8)
    scan_x = np.concatenate(scan_x_v, dtype=np.float32)
    scan_y = np.concatenate(scan_y_v, dtype=np.float32)
    scan_len = np.concatenate(scan_len_v, dtype=np.int16)
    label = np.concatenate(label_v, dtype=np.int8)

    assert sum(scan_len) == len(scan_x) == len(scan_y)
    print(f'最终pos: {sum(label == 1)}, 最终neg：{sum(label == 0)}')

    np.savez(dir_npz,
             pr_mz=pr_mz,
             pr_charge=pr_charge,
             scan_x=scan_x,
             scan_y=scan_y,
             scan_len=scan_len,
             label=label)


if __name__ == '__main__':
    main()
    print('finished.')